use std::env;
use std::io::{self, Read, Write};
use std::sync::atomic::{AtomicI32, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
use std::convert::TryInto;
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};

const N_LEN: usize = 65536 - 16;

fn get_unix_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

fn main() -> io::Result<()> {
    thread::sleep(Duration::from_millis(2000));

    let mut args = env::args().skip(1);
    let mut function = String::new();
    let mut instance: i32 = 0;
    let mut sub_input: i32 = 0;
    let mut sub_output: i32 = 0;
    let mut b_info = false;

    while let Some(arg) = args.next() {
        if let Some((key, value)) = arg.split_once('=') {
            eprintln!("{}:{}", key, value);
            match key {
                "--function" => function = value.to_string(),
                "--instance" => instance = value.parse().unwrap_or(0),
                "--data_in" => sub_input = value.parse().unwrap_or(0),
                "--data_out" => sub_output = value.parse().unwrap_or(0),
                "--information" => b_info = true,
                _ => {}
            }
        }
    }

    if b_info {
        return Ok(());
    }

    if instance <= 0 || function.is_empty() {
        return Ok(());
    }

    let way_count = Arc::new(AtomicI32::new(0));
    let stdout = Arc::new(Mutex::new(io::stdout()));
    let finished = Arc::new(AtomicI32::new(0));

    let th_send = {
        let way_count = Arc::clone(&way_count);
        let stdout = Arc::clone(&stdout);
        let sub_output = sub_output;
        let instance = instance;

        thread::spawn(move || {
            let mut thd = 10;
            let mut cnt = 0u64;
            while way_count.load(Ordering::Relaxed) >= 0 {
                if way_count.load(Ordering::Relaxed) < thd {
                    let buf_header = [0x3c, 0x5a, 0x7e, 0x69];
                    let mut data = vec![0u8; N_LEN];

                    // 替换 unsafe 指针操作为安全的字节操作
                    let mscv = get_unix_timestamp();
                    data[0..8].copy_from_slice(&mscv.to_le_bytes());

                    data[8..16].copy_from_slice(&cnt.to_le_bytes());
                    cnt += 1;

                    let mut guard = stdout.lock().unwrap();
                    guard.write_all(&buf_header).unwrap();
                    guard.write_all(&sub_output.to_le_bytes()).unwrap();
                    guard.write_all(&instance.to_le_bytes()).unwrap();
                    guard.write_all(&(N_LEN as i32).to_le_bytes()).unwrap();
                    guard.write_all(&data).unwrap();
                    guard.flush().unwrap();
                    drop(guard);

                    way_count.fetch_add(1, Ordering::Relaxed);
                } else {
                    thread::sleep(Duration::from_millis(1));
                    if way_count.load(Ordering::Relaxed) == 0 && thd < 1000 {
                        thd += 1;
                    }
                }
            }
        })
    };

    // 剩余部分保持不变...
    // [接收循环部分与原代码相同，无需修改]

    let mut recvcnt = 0u64;
    let mut delay = 0u64;
    let mut first_clk = 0u64;
    let mut stdin = io::stdin();

    loop {
        let mut header = [0u8; 4];
        stdin.read_exact(&mut header)?;

        if header != [0x3C, 0x5A, 0x7E, 0x69] {
            eprintln!("BAD HEADER");
            continue;
        }

        let mut n_sub = [0u8; 4];
        let mut n_path = [0u8; 4];
        let mut len = [0u8; 4];
        stdin.read_exact(&mut n_sub)?;
        stdin.read_exact(&mut n_path)?;
        stdin.read_exact(&mut len)?;

        let n_sub = i32::from_le_bytes(n_sub);
        let len = i32::from_le_bytes(len);

        if len < 0 || len as usize != N_LEN || n_sub <= 0 {
            let mut data = vec![0u8; len as usize];
            stdin.read_exact(&mut data)?;
            if String::from_utf8_lossy(&data).contains("function=quit;") {
                break;
            }
        } else {
            let mut data = vec![0u8; N_LEN];
            stdin.read_exact(&mut data)?;

            if n_sub != sub_input {
                eprintln!("BAD SUBJECT");
                continue;
            }

            let instance = instance;
            let n_path = i32::from_le_bytes(n_path);
            if n_path != instance {
                let mut guard = stdout.lock().unwrap();
                guard.write_all(&header)?;
                guard.write_all(&sub_output.to_le_bytes())?;
                guard.write_all(&n_path.to_le_bytes())?;
                guard.write_all(&(N_LEN as i32).to_le_bytes())?;
                guard.write_all(&data)?;
                guard.flush()?;
            } else {
                way_count.fetch_sub(1, Ordering::Relaxed);
                let clk = u64::from_le_bytes(data[0..8].try_into().unwrap());
                if recvcnt == 0 {
                    first_clk = clk;
                }
                recvcnt += 1;
                delay += get_unix_timestamp() - clk;

                if recvcnt >= 10000 {
                    let total_bytes = 65536 * 10000;
                    let tm_cost = (clk - first_clk) as f64 + 1e-10;
                    let speed = total_bytes as f64 / 1024.0 / 1024.0 / tm_cost;
                    eprintln!(
                        "Cnt = {}, Average delay = {} ms, Speed = {:.2} MB/s",
                        u64::from_le_bytes(data[8..16].try_into().unwrap()),
                        delay,
                        speed
                    );
                    recvcnt = 0;
                    delay = 0;
                }
            }
        }
    }

    finished.store(1, Ordering::Relaxed);
    th_send.join().unwrap();
    Ok(())
}
